import os
from multiprocessing import Pool
from scipy.signal import resample, correlate
from scipy.io import loadmat, savemat
import numpy as np
import time
from sklearn.utils import shuffle
import scipy
from scikits import audiolab
import glob


def generate_lst():
    root_path = '/home/xiaomin/wxm/Data/KaggleSeizure2016/RawData'
    train_lst = 'train.lst'
    test_lst = 'test.lst'
    f_train_lst = open(root_path + '/' + train_lst, 'w')
    f_test_lst = open(root_path + '/' + test_lst, 'w')
    nb_train = 0
    nb_train_0 = 0
    nb_train_1 = 0
    nb_test = 0
    for i in os.listdir(root_path):
        if (not i.endswith('.zip')) and (not i.endswith('.csv')) and (not i.endswith('.lst')) and (
                not i.endswith('.npy')):
            if 'train' in i:
                print 'train:' + i
                for j in os.listdir(root_path + '/' + i):
                    line = str(nb_train) + '\t' + i + '/' + j + '\t' + j[-5:-4] + '\n'
                    f_train_lst.write(line)
                    if int(j[-5:-4]) == 0:
                        nb_train_0 += 1
                    elif int(j[-5:-4]) == 1:
                        nb_train_1 += 1
                    nb_train += 1
            elif 'test' in i:
                print 'test: ' + i
                for j in os.listdir(root_path + '/' + i):
                    line = str(nb_test) + '\t' + i + '/' + j + '\n'
                    f_test_lst.write(line)
                    nb_test += 1
    print nb_train_1
    print nb_train_0
    print nb_train
    print nb_test


def generate_data_by_Fourier_resample(train_lst, test_lst, correlate_flag=False, Debug=False):
    root_dir = '/home/xiaomin/wxm/Data/KaggleSeizure2016/RawData'
    npy_dir = '/home/xiaomin/wxm/Data/KaggleSeizure2016/npys'
    folder = 'resample600w400_NoCorr'
    period = 300
    if not os.path.exists(npy_dir + '/' + folder):
        os.makedirs(npy_dir + '/' + folder)
    f_train = open(root_dir + '/' + train_lst, 'r')
    f_test = open(root_dir + '/' + test_lst, 'r')
    num_train = len(open(root_dir + '/' + train_lst, 'r').readlines())
    num_test = len(open(root_dir + '/' + test_lst, 'r').readlines())

    train_x = []
    train_y = []
    test_x = []
    nb_train = 0
    nb_test = 0

    for i in range(num_train):
        print 'train'
        nb_train += 1
        if Debug is True and nb_train >= 50:
            break
        print nb_train
        index, x_path, label = f_train.readline().strip('\n').split('\t')
        print x_path
        if '1_45_1' in x_path:
            print '*' * 50
            continue
        mat = loadmat(root_dir + '/' + x_path)
        print mat.keys()
        mat_arr = mat['dataStruct'][0][0][0].transpose()
        print mat_arr.shape
        mat_arr_resample = resample(mat_arr, 600, axis=1, window=400)
        print mat_arr_resample.shape
        correlates = None
        if correlate_flag is True:
            for m in range(len(mat_arr)):
                correlate_m = np.array([])
                for n in range(len(mat_arr)):
                    if m != n:
                        corr_m = correlate(mat_arr_resample[m], mat_arr_resample[n], mode='same')
                        correlate_m = np.concatenate([correlate_m, corr_m])
                if correlates is None:
                    correlates = correlate_m
                elif correlates is not None:
                    correlates = np.vstack([correlates, correlate_m])
        if correlate_flag is True:
            x = np.column_stack([mat_arr_resample, correlates])
        elif correlate_flag is False:
            x = mat_arr_resample

        train_x.append(x)
        train_y.append(int(label))
        print '*' * 50
        if nb_train % period == 0:
            train_x = np.array(train_x)
            train_y = np.array(train_y)
            print train_x.shape
            print train_y.shape
            np.save(npy_dir + '/' + folder + '/' + 'train_x_' + str(nb_train) + '.npy', train_x)
            np.save(npy_dir + '/' + folder + '/' + 'train_y_' + str(nb_train) + '.npy', train_y)
            train_x = []
            train_y = []

    for i in range(num_test):
        print 'test'
        nb_test += 1
        if Debug is True and nb_test >= 50:
            break
        print nb_test
        index, y_path = f_test.readline().strip('\n').split('\t')
        print y_path
        mat = loadmat(root_dir + '/' + y_path)
        print mat.keys()
        mat_arr = mat['dataStruct'][0][0][0].transpose()
        print mat_arr.shape
        mat_arr_resample = resample(mat_arr, 600, axis=1, window=400)
        print mat_arr_resample.shape
        correlates = None
        if correlate_flag is True:
            for m in range(len(mat_arr)):
                correlate_m = np.array([])
                for n in range(len(mat_arr)):
                    if m != n:
                        corr_m = correlate(mat_arr_resample[m], mat_arr_resample[n], mode='same')
                        correlate_m = np.concatenate([correlate_m, corr_m])
                if correlates is None:
                    correlates = correlate_m
                elif correlates is not None:
                    correlates = np.vstack([correlates, correlate_m])
        if correlate_flag is True:
            x = np.column_stack([mat_arr_resample, correlates])
        elif correlate_flag is False:
            x = mat_arr_resample
        test_x.append(x)
        print '*' * 50
        if nb_test % period == 0:
            test_x = np.array(test_x)
            np.save(npy_dir + '/' + folder + '/' + 'test_x_' + str(nb_test) + '.npy', test_x)
            print test_x.shape
            test_x = []


def generate_data_by_Fourier_resample_mulyiprocess(input_index):
    lst = 'test.lst'
    folder = 'resample600w400_Corr'
    Debug = False
    correlate_flag = True
    root_dir = '/home/xiaomin/wxm/Data/KaggleSeizure2016/RawData'
    mat_dir = '/home/xiaomin/wxm/Data/KaggleSeizure2016/mats'
    if not os.path.exists(mat_dir + '/' + folder):
        os.makedirs(mat_dir + '/' + folder)
    f = open(root_dir + '/' + lst, 'r')
    num = len(open(root_dir + '/' + lst, 'r').readlines())

    x = None
    kv_train_x = {}
    kv_train_y = {}
    kv_test_x = {}

    if 'train' in lst:
        print 'train'
        for i in range(num):
            if Debug is True and input_index >= 50:
                break
            index, x_path, label = f.readline().strip('\n').split('\t')
            kv_train_x[str(index)] = x_path
            kv_train_y[str(index)] = str(label)
        print input_index
        if 'train_1/1_45_1.mat' == kv_train_x[str(input_index)]:
            input_index += 1
        mat = loadmat(root_dir + '/' + kv_train_x[str(input_index)])
        print mat.keys()
        mat_arr = mat['dataStruct'][0][0][0].transpose()
        print mat_arr.shape
        mat_arr_resample = resample(mat_arr, 600, axis=1, window=400)
        print mat_arr_resample.shape
        correlates = None
        if correlate_flag is True:
            for m in range(len(mat_arr)):
                correlate_m = np.array([])
                for n in range(len(mat_arr)):
                    if m != n:
                        corr_m = correlate(mat_arr_resample[m], mat_arr_resample[n], mode='same')
                        correlate_m = np.concatenate([correlate_m, corr_m])
                if correlates is None:
                    correlates = correlate_m
                elif correlates is not None:
                    correlates = np.vstack([correlates, correlate_m])
        if correlate_flag is True:
            x = np.column_stack([mat_arr_resample, correlates])
        elif correlate_flag is False:
            x = mat_arr_resample

        mat_save = {'x': x,
                    'y': int(kv_train_y[str(input_index)])}
        savemat(mat_dir + '/' + folder + '/' + 'train_' + str(input_index) + '.mat', mat_save)
        print '*' * 50

    elif 'test' in lst:
        print 'test'
        for i in range(num):
            if Debug is True and input_index >= 50:
                break
            index, x_path = f.readline().strip('\n').split('\t')
            kv_test_x[str(index)] = x_path
        print input_index
        mat = loadmat(root_dir + '/' + kv_test_x[str(input_index)])
        print mat.keys()
        mat_arr = mat['dataStruct'][0][0][0].transpose()
        print mat_arr.shape
        mat_arr_resample = resample(mat_arr, 600, axis=1, window=400)
        print mat_arr_resample.shape
        correlates = None
        if correlate_flag is True:
            for m in range(len(mat_arr)):
                correlate_m = np.array([])
                for n in range(len(mat_arr)):
                    if m != n:
                        corr_m = correlate(mat_arr_resample[m], mat_arr_resample[n], mode='same')
                        correlate_m = np.concatenate([correlate_m, corr_m])
                if correlates is None:
                    correlates = correlate_m
                elif correlates is not None:
                    correlates = np.vstack([correlates, correlate_m])
        if correlate_flag is True:
            x = np.column_stack([mat_arr_resample, correlates])
        elif correlate_flag is False:
            x = mat_arr_resample
        mat_save = {'x': x}
        savemat(mat_dir + '/' + folder + '/' + 'test_' + str(input_index) + '.mat', mat_save)
        print '*' * 50


def generate_data_period_npy(folder):
    root_dir = '/home/xiaomin/wxm/Data/KaggleSeizure2016/mats'
    npy_dir = '/home/xiaomin/wxm/Data/KaggleSeizure2016/npy'
    if not os.path.exists(npy_dir + '/' + folder):
        os.makedirs(npy_dir + '/' + folder)
    train_x = None
    train_y = None
    test_x = None
    nb_mat = 0
    nb_train = 0
    nb_test = 0
    period = 500
    for mat in os.listdir(root_dir + '/' + folder):
        begin = time.time()
        nb_mat += 1
        print mat
        print nb_mat
        if 'train' in mat:
            nb_train += 1
            if nb_train % period == 0:
                np.save(npy_dir + '/' + folder + '/' + 'train_' + str(nb_train) + '.npy')
            arr = loadmat(root_dir + '/' + folder + '/' + mat)
            x = np.expand_dims(arr['x'], 0)
            print x.shape
            if train_x is None:
                train_x = x
            elif train_x is not None:
                train_x = np.concatenate([train_x, x], 0)

            y = np.expand_dims(arr['y'], 0)
            if train_y is None:
                train_y = y
            elif train_y is not None:
                train_y = np.concatenate([train_y, y], 0)

        elif 'test' in mat:
            nb_test += 1
            if nb_test % period == 0:
                np.save(npy_dir + '/' + folder + '/' + 'train_' + str(nb_test) + '.npy')
            arr = loadmat(root_dir + '/' + folder + '/' + mat)
            x = np.expand_dims(arr['x'], 0)
            if test_x is None:
                test_x = x
            elif test_x is not None:
                test_x = np.concatenate([test_x, x], 0)

        end = time.time()
        print 'time per mat: ' + str(end - begin)
        print '*' * 50

    np.save(npy_dir + '/' + folder + '/' + 'train_' + str(nb_train) + '.npy')
    np.save(npy_dir + '/' + folder + '/' + 'train_' + str(nb_test) + '.npy')
    print 'train_x.shape: ' + str(train_x.shape)
    print 'train_y.shape: ' + str(train_y.shape)
    print 'test_x.shape: ' + str(test_x.shape)
    return train_x, train_y, test_x


def generate_train_val_list():
    root_path = '/home/xiaomin/wxm/Data/KaggleSeizure2016/wavs'
    trainall_lst = 'rawtrain_allsubjs_e0.lst'
    train_lst = '/home/xiaomin/wxm/Data/KaggleSeizure2016/wavs/train_allsubjs_e0.lst'
    val_lst = '/home/xiaomin/wxm/Data/KaggleSeizure2016/wavs/val_allsubjs_e0.lst'

    f = open(root_path + '/' + trainall_lst, 'r')
    f_train = open(train_lst, 'w')
    f_val = open(val_lst, 'w')
    num = len(open(root_path + '/' + trainall_lst, 'r').readlines())
    lines = f.readlines()
    nb = 0
    lines = shuffle(lines)
    for i in range(num):
        nb += 1
        if nb <= 5000:
            f_train.write(lines[i])
        elif nb > 5000:
            f_val.write(lines[i])


def generate_separate_channel_wav_from_mat_per_subj_multiprocessing(index=0):
    subj = 'test_3'
    root_folder = '/home/xiaomin/wxm/Data/KaggleSeizure2016/RawData'
    wav_folder = '/home/xiaomin/wxm/Data/KaggleSeizure2016/wavs/SubjPerChannel'
    data_folder = root_folder + '/' + subj
    save_folder = wav_folder + '/' + subj
    # if not os.path.exists(save_folder):
        # os.makedirs(save_folder)

    # mat_paths = []
    # for mat_name in os.listdir(data_folder):
    #     if mat_name.endswith('.mat'):
    #         mat_path = data_folder + '/' + mat_name
    #         mat_paths.append(mat_path)
    mat_paths = glob.glob(data_folder + '/' + '*.mat')
    print str(index) + 'th in ' + str(len(mat_paths)) + '\t' + mat_paths[index] + '\n'
    splits = mat_paths[index].split('/')
    mat = scipy.io.loadmat(mat_paths[index])
    dat = mat['dataStruct'][0][0][0]

    mn = dat.min()
    mx = dat.max()
    mx = float(max(abs(mx), abs(mn)))
    if mx != 0:
        dat *= 0x7FFF / mx
    dat = np.int16(dat)
    for elec in range(16):
        dstfile = save_folder + '/' + splits[-1][0:-4] + '_e' + str(elec) + '.wav'
        aud = dat[:, elec]
        audiolab.wavwrite(aud, dstfile, fs=400, enc='pcm16')


def generate_lst_from_exit_lst():
    exit_lst_path = '/home/xiaomin/wxm/Data/KaggleSeizure2016/mats/val.lst'
    res_lst_path = '/home/xiaomin/wxm/Data/KaggleSeizure2016/wavs/val_allsubjs_e0.lst'

    num = len(open(exit_lst_path, 'r').readlines())
    f_exit = open(exit_lst_path, 'r')
    f_res = open(res_lst_path, 'w')

    for i in range(num):
        print i
        line = f_exit.readline()
        index, path, label = line.strip('\n').split('\t')
        new_path = path[0:-4] + '_e0.wav'
        new_line = index + '\t' + new_path + '\t' + label + '\n'
        f_res.write(new_line)


def generate_lst_from_exit_folders():
    exit_root_folder = '/home/xiaomin/wxm/Data/KaggleSeizure2016/wavs/SubjPerChannel'
    res_lst_path = '/home/xiaomin/wxm/Data/KaggleSeizure2016/wavs/rawtrain_allsubjs_e0.lst'

    nb_wav = 0
    f_res = open(res_lst_path, 'w')
    for sub_folder in os.listdir(exit_root_folder):
        if sub_folder.startswith('train'):
            sub_folder_path = exit_root_folder + '/' + sub_folder
            for wav_name in os.listdir(sub_folder_path):
                if '_e0' in wav_name:
                    print nb_wav
                    print wav_name
                    splits = wav_name.split('_')
                    label = splits[2]
                    line = str(nb_wav) + '\t' + sub_folder + '/' + wav_name + '\t' + str(label) + '\n'
                    f_res.write(line)
                    nb_wav += 1
                    print '*' * 50


if __name__ == '__main__':
    # p = Pool(18)
    # p.map(generate_data_by_Fourier_resample_mulyiprocess, range(0, 6042))
    # generate_data_by_Fourier_resample_mulyiprocess(84)
    # generate_lst()
    # generate_data_period_npy('resample600w400_Corr')
    # generate_train_val_list()
    p = Pool(16)
    p.map(generate_separate_channel_wav_from_mat_per_subj_multiprocessing, range(0, 2286))
    # generate_separate_channel_wav_from_mat_per_subj_multiprocessing()
    # generate_lst_from_exit_lst()
    # generate_train_val_list()
    # generate_lst_from_exit_folders()
